import os
from urllib.parse import urlparse
from fastmcp import FastMCP
from fastmcp.server.dependencies import get_http_request
from fastmcp.server.dependencies import get_http_headers
import base64
from fastmcp.server.context import Context
from typing import Annotated
from pydantic import Field
from jimeng_client import JiMengClient
import logging

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

mcp = FastMCP("image-gen-server")
client = JiMengClient("AKLTNDVmYWYxNjA1YTE4NDQ3NTk4MTc5NDJhNjI2NzUzMDA",
                      "WW1ZNU16QTBPRFZrTTJRMk5EZGpZV0ptTlRRek1XRXdNVGN6WWpBMU4yUQ==")


@mcp.tool(
    name="generate_image",
    description="根据文本描述生成图片，并以base64格式返回图片"
)
async def generate_image(
        prompt: Annotated[str, Field(description="图片的文本prompt描述,500字符长度限制，一个汉字算一个字符长度", max_length=500)],
        filepath: Annotated[str, Field(description="图片保存路径（本地绝对路径）")],
        width: Annotated[int, Field(
            description="生成图片的宽度，单位：像素", ge=512, le=768)] = 512,
        height:  Annotated[int, Field(
            description="生成图片的高度，单位：像素", ge=512, le=768)] = 512
) -> dict:
    """
    根据文本描述生成图片，并以base64格式返回图片
    """
    logger.info(f"generate_image: {prompt}")

    headers = get_http_headers()
    client.access_key = headers.get("ak")
    client.secret_key = headers.get("sk")

    response = client.generate_image(prompt, width, height)
    if response and response.is_success:
        save_image_response = save_image(response.image, filepath)
        save_image_response["filepath"] = filepath
        return save_image_response
    else:
        return {
            "filepath": None,
            "message": f"error: {response.message}"
        }



def save_image(
        image: bytes,
        filepath: str
) -> dict:
    logger.info(f"save image to {filepath}")
    if os.path.exists(filepath):
        return {
            "message": "error: file already exists"
        }
    if os.path.isdir(filepath):
        return {
            "message": "error: filepath is a directory"
        }
    filedir = os.path.dirname(filepath)
    if not os.path.exists(filedir):
        os.makedirs(filedir)

    with open(filepath, "wb") as f:
        f.write(image)

    return {
        "message": "success",
        "filepath": filepath
    }


if __name__ == '__main__':
    logger.info("启动图像生成服务...")
    mcp.run(transport="streamable-http", host="127.0.0.1",
            path="/jimeng_mcp", port=9000, log_level="info")
